#!/usr/bin/env python3
from rdkit import Chem
import pandas as pd
import argparse

def inchi_to_canonical_smiles(inchi_string):
    mol = Chem.MolFromInchi(inchi_string, sanitize=False)
    if mol is None:
        raise ValueError("Invalid InChI string provided.")
    
    try:
        Chem.SanitizeMol(mol)
    except:
        print(f'WARNING: fail to sanitize {inchi_string}')
        pass
    
    canonical_smiles = Chem.MolToSmiles(mol, canonical=True)
    return canonical_smiles


def main():
    parser = argparse.ArgumentParser(description="Convert InChI to rdkit canonical SMILES.")
    parser.add_argument("-i", "--input", required=True, help="Input TSV/CSV file with InChI column")
    parser.add_argument("-o", "--output", required=True, help="Output smi file")
    parser.add_argument("-f", "--failed", default="failed_ids.txt", help="Failed IDs output file")
    parser.add_argument("-c", "--concurrency", type=int, default=10, help="Number of concurrent downloads")
    args = parser.parse_args()

    try:
        df = pd.read_csv(args.input, sep=None, engine="python")
        if "InChI" not in df.columns or "zinc_id" not in df.columns:
            raise ValueError("Input TSV must contain a 'InChI' column & a 'zinc_id' column.")
        inchis = df["InChI"]
        print(f"Number of InChI: {len(inchis)}")
    except Exception as e:
        print(f"Error reading input file: {e}")
        return
    
    canon_smi = []

    for inchi in inchis:
        canon_smi.append(inchi_to_canonical_smiles(inchi))
    
    df["can_smiles"] = canon_smi
    
    df[["can_smiles","zinc_id"]].to_csv(args.output, sep=" ", index=False, header=False)

if __name__ == "__main__":
    main()
